"""
文件名: security.py
描述: 安全工具模块

本模块包含以下主要功能：
1. JWT令牌的创建和验证
2. 密码哈希和验证
3. OAuth令牌的加密和解密存储
4. 安全相关的辅助函数
5. 加密密钥管理

依赖模块:
   - jose: JWT处理
   - passlib: 密码哈希
   - cryptography: 加密解密
   - datetime: 时间处理
   - typing: 类型注解

使用示例:
   >>> from src.utils.security import create_access_token, verify_token
   >>> token = create_access_token({"sub": "user123"})
   >>> payload = verify_token(token)
   >>> 
   >>> # 加密存储OAuth令牌
   >>> encrypted = encrypt_token("oauth_access_token")
   >>> decrypted = decrypt_token(encrypted)

注意事项:
   - JWT密钥必须保密且足够复杂
   - 密码哈希使用bcrypt算法
   - OAuth令牌加密存储到数据库
   - 定期更换加密密钥
   - 遵循安全最佳实践

作者: AI助手
创建日期: 2024-06-27
最后修改: 2024-06-27
版本: 1.0.0
"""

import base64
import secrets
from datetime import datetime, timedelta
from typing import Any, Dict, Optional, Union

from cryptography.fernet import Fernet
from jose import JWTError, jwt
from passlib.context import CryptContext
from passlib.hash import bcrypt

from src.core.config import get_settings
from src.core.logger import get_logger

logger = get_logger(__name__)
settings = get_settings()

# 密码哈希上下文
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# 加密密钥（用于OAuth令牌加密）
# 注意：生产环境中应该从环境变量或密钥管理服务获取
ENCRYPTION_KEY = base64.urlsafe_b64encode(settings.secret_key.encode()[:32].ljust(32, b'0'))
fernet = Fernet(ENCRYPTION_KEY)


def create_access_token(
    data: Dict[str, Any], 
    expires_delta: Optional[timedelta] = None
) -> str:
    """
    创建JWT访问令牌
    
    Args:
        data: 要编码到令牌中的数据
        expires_delta: 令牌有效期，默认使用配置文件中的设置
        
    Returns:
        str: JWT令牌字符串
        
    Raises:
        Exception: 令牌创建失败时抛出异常
    """
    try:
        to_encode = data.copy()
        
        # 设置过期时间
        if expires_delta:
            expire = datetime.utcnow() + expires_delta
        else:
            expire = datetime.utcnow() + timedelta(
                minutes=settings.access_token_expire_minutes
            )
        
        to_encode.update({
            "exp": expire,
            "iat": datetime.utcnow(),
            "iss": settings.app_name,  # 发行者
            "type": "access_token"  # 令牌类型
        })
        
        # 创建JWT令牌
        encoded_jwt = jwt.encode(
            to_encode, 
            settings.secret_key, 
            algorithm=settings.algorithm
        )
        
        logger.debug("JWT访问令牌创建成功", subject=data.get("sub"))
        return encoded_jwt
        
    except Exception as e:
        logger.error(f"JWT令牌创建失败: {e}", error=str(e), data=data)
        raise


def create_refresh_token(data: Dict[str, Any]) -> str:
    """
    创建JWT刷新令牌
    
    Args:
        data: 要编码到令牌中的数据
        
    Returns:
        str: JWT刷新令牌字符串
    """
    try:
        to_encode = data.copy()
        
        # 刷新令牌有效期更长（7天）
        expire = datetime.utcnow() + timedelta(days=7)
        
        to_encode.update({
            "exp": expire,
            "iat": datetime.utcnow(),
            "iss": settings.app_name,
            "type": "refresh_token"
        })
        
        encoded_jwt = jwt.encode(
            to_encode,
            settings.secret_key,
            algorithm=settings.algorithm
        )
        
        logger.debug("JWT刷新令牌创建成功", subject=data.get("sub"))
        return encoded_jwt
        
    except Exception as e:
        logger.error(f"JWT刷新令牌创建失败: {e}", error=str(e), data=data)
        raise


def verify_token(token: str, token_type: str = "access_token") -> Optional[Dict[str, Any]]:
    """
    验证JWT令牌
    
    Args:
        token: JWT令牌字符串
        token_type: 令牌类型（access_token 或 refresh_token）
        
    Returns:
        Optional[Dict[str, Any]]: 解码后的令牌数据，验证失败返回None
    """
    try:
        # 解码JWT令牌
        payload = jwt.decode(
            token,
            settings.secret_key,
            algorithms=[settings.algorithm]
        )
        
        # 验证令牌类型
        if payload.get("type") != token_type:
            logger.warning(f"令牌类型不匹配: 期望 {token_type}, 实际 {payload.get('type')}")
            return None
        
        # 验证发行者
        if payload.get("iss") != settings.app_name:
            logger.warning(f"令牌发行者不匹配: 期望 {settings.app_name}, 实际 {payload.get('iss')}")
            return None
        
        # 检查用户ID
        user_id = payload.get("sub")
        if not user_id:
            logger.warning("令牌中缺少用户ID")
            return None
        
        logger.debug("JWT令牌验证成功", user_id=user_id, token_type=token_type)
        return payload
        
    except JWTError as e:
        logger.warning(f"JWT令牌验证失败: {e}", error=str(e))
        return None
    except Exception as e:
        logger.error(f"令牌验证过程中发生错误: {e}", error=str(e))
        return None


def get_password_hash(password: str) -> str:
    """
    生成密码哈希
    
    Args:
        password: 明文密码
        
    Returns:
        str: 哈希后的密码
    """
    try:
        hashed = pwd_context.hash(password)
        logger.debug("密码哈希生成成功")
        return hashed
    except Exception as e:
        logger.error(f"密码哈希生成失败: {e}", error=str(e))
        raise


def verify_password(plain_password: str, hashed_password: str) -> bool:
    """
    验证密码
    
    Args:
        plain_password: 明文密码
        hashed_password: 哈希密码
        
    Returns:
        bool: 密码是否正确
    """
    try:
        result = pwd_context.verify(plain_password, hashed_password)
        logger.debug("密码验证完成", result=result)
        return result
    except Exception as e:
        logger.error(f"密码验证失败: {e}", error=str(e))
        return False


def encrypt_token(token: str) -> str:
    """
    加密OAuth令牌用于数据库存储
    
    Args:
        token: 明文令牌
        
    Returns:
        str: 加密后的令牌（Base64编码）
        
    Raises:
        Exception: 加密失败时抛出异常
    """
    try:
        if not token:
            return ""
        
        # 加密令牌
        encrypted_token = fernet.encrypt(token.encode('utf-8'))
        
        # Base64编码用于数据库存储
        encoded_token = base64.urlsafe_b64encode(encrypted_token).decode('utf-8')
        
        logger.debug("令牌加密成功")
        return encoded_token
        
    except Exception as e:
        logger.error(f"令牌加密失败: {e}", error=str(e))
        raise


def decrypt_token(encrypted_token: str) -> str:
    """
    解密OAuth令牌
    
    Args:
        encrypted_token: 加密的令牌（Base64编码）
        
    Returns:
        str: 解密后的明文令牌
        
    Raises:
        Exception: 解密失败时抛出异常
    """
    try:
        if not encrypted_token:
            return ""
        
        # Base64解码
        encrypted_data = base64.urlsafe_b64decode(encrypted_token.encode('utf-8'))
        
        # 解密令牌
        decrypted_token = fernet.decrypt(encrypted_data).decode('utf-8')
        
        logger.debug("令牌解密成功")
        return decrypted_token
        
    except Exception as e:
        logger.error(f"令牌解密失败: {e}", error=str(e))
        raise


def generate_random_key(length: int = 32) -> str:
    """
    生成随机密钥
    
    Args:
        length: 密钥长度（字节）
        
    Returns:
        str: 随机密钥（Base64编码）
    """
    try:
        random_bytes = secrets.token_bytes(length)
        key = base64.urlsafe_b64encode(random_bytes).decode('utf-8')
        logger.debug(f"随机密钥生成成功，长度: {length}")
        return key
    except Exception as e:
        logger.error(f"随机密钥生成失败: {e}", error=str(e))
        raise


def generate_state_token() -> str:
    """
    生成OAuth状态令牌（防CSRF攻击）
    
    Returns:
        str: 状态令牌
    """
    try:
        state = secrets.token_urlsafe(32)
        logger.debug("OAuth状态令牌生成成功")
        return state
    except Exception as e:
        logger.error(f"OAuth状态令牌生成失败: {e}", error=str(e))
        raise


def validate_token_format(token: str) -> bool:
    """
    验证令牌格式是否正确
    
    Args:
        token: 要验证的令牌
        
    Returns:
        bool: 格式是否正确
    """
    try:
        if not token or not isinstance(token, str):
            return False
        
        # JWT令牌应该有三个部分，用.分隔
        parts = token.split('.')
        if len(parts) != 3:
            return False
        
        # 检查每个部分是否为有效的Base64编码
        for part in parts:
            try:
                # 添加必要的填充
                padded = part + '=' * (4 - len(part) % 4)
                base64.urlsafe_b64decode(padded)
            except Exception:
                return False
        
        return True
        
    except Exception as e:
        logger.error(f"令牌格式验证失败: {e}", error=str(e))
        return False


# 令牌相关常量
TOKEN_TYPE_ACCESS = "access_token"
TOKEN_TYPE_REFRESH = "refresh_token"
TOKEN_TYPE_OAUTH = "oauth_token"

# 导出主要接口
__all__ = [
    "create_access_token",
    "create_refresh_token",
    "verify_token",
    "get_password_hash",
    "verify_password",
    "encrypt_token",
    "decrypt_token",
    "generate_random_key",
    "generate_state_token",
    "validate_token_format",
    "TOKEN_TYPE_ACCESS",
    "TOKEN_TYPE_REFRESH",
    "TOKEN_TYPE_OAUTH",
]